package xinbloom

import (
	"crypto/sha256"
	"encoding/binary"
	"hash/fnv"
	"math"
)

type HashFunc func([]byte) uint64

type BloomFilter struct {
	bm        *bitmap    // 存储数据
	hashFuncs []HashFunc // 一组 hash 函数
}

// NewBloom 初始化布隆过滤器
// n 预期数据量
// p 误判率
// m 存储空间大小
// k hash 函数个数
func NewBloom(n int, p float64) *BloomFilter {
	// m k 的计算方式，参考 Java Guava 库的布隆过滤器实现。

	// 计算最佳 m 值
	m := -float64(n) * math.Log(p) / (math.Log(2) * math.Log(2))

	// 计算最佳 k 值
	// 一般都是布隆过滤器自行实现
	// k := math.Max(1, math.Round(m/float64(n)*math.Log(2)))

	hashFuncs := []HashFunc{
		fnv_1, sha256_1,
	}

	return &BloomFilter{
		bm:        newBitMap(int(m)),
		hashFuncs: hashFuncs,
	}
}

// Check 验证元素是否存在
func (bf *BloomFilter) Check(item string) bool {
	for _, f := range bf.hashFuncs {
		hashcode := f([]byte(item))

		index := hashcode % bf.bm.size()

		// 判断 bitmap 中是否存在
		flag := bf.bm.test(index)
		if !flag {
			return false
		}
	}

	return true
}

// Add 添加元素
func (bf *BloomFilter) Add(item string) {
	for _, f := range bf.hashFuncs {
		hashcode := f([]byte(item))

		index := hashcode % bf.bm.size()

		// 判断 bitmap 中是否存在
		flag := bf.bm.test(index)
		if !flag {
			bf.bm.set(index)
		}
	}
}

// AddHashFunc 新增 hash 函数
func (bf *BloomFilter) AddHashFunc(f HashFunc) {
	bf.hashFuncs = append(bf.hashFuncs, f)
}

func fnv_1(data []byte) uint64 {
	h := fnv.New64a()
	h.Write(data)
	return h.Sum64()
}

func sha256_1(data []byte) uint64 {
	h := sha256.New()
	h.Write(data)
	return binary.BigEndian.Uint64(h.Sum(nil)[:8])
}
